#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelHistc
{
public:
    __aicore__ inline KernelHistc() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t bins, float min, float max,
                                uint32_t blockLength, uint32_t tileNum, uint32_t tileLength,
                                uint32_t lasttileLength)
    {

        this->bins = bins;
        this->min = min;
        this->max = max;
        this->blockLength = blockLength;
        this->tileNum = tileNum;
        this->lasttileLength = lasttileLength;
        this->tileLength = tileLength / BUFFER_NUM;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, this->tileLength * tileNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->bins);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        this->spareNum = xGm.GetValue(0);
        if (min == 0 && min == max)
        {
            this->usedataValue = 1;
            this->min = this->spareNum;
            this->max = this->spareNum;
        }
        else
        {
            this->usedataValue = 0;
        }

        pipe.InitBuffer(calcXF, this->tileLength * sizeof(float));
        pipe.InitBuffer(Tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(Tmp2, this->tileLength * sizeof(uint8_t));
        pipe.InitBuffer(Tmp3, this->tileLength * sizeof(int32_t));
        pipe.InitBuffer(Tmp4, this->tileLength * sizeof(float));
        pipe.InitBuffer(calcZero, this->tileLength * sizeof(float));
        pipe.InitBuffer(calcTensor, this->tileLength * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;

        if (this->usedataValue == 1)
        {
            for (int32_t i = 0; i < loopCount; i++)
            {
                CopyIn2(i);
                SetMM(i);
            }
        }

        float RemainNum = (BUFFER_NUM * this->tileLength - this->lasttileLength);

        for (int32_t i = 0; i < loopCount; i++)
        {
            CopyIn(i);
            Compute(i);
        }
        int32_t Error = (-this->min * this->bins) / (this->max - this->min);
        if (Error >= 0 && Error < this->bins)
        {
            float Tvalue = (float)yGm.GetValue(Error) - (float)RemainNum;
            yGm.SetValue(Error, (DTYPE_X)Tvalue); // fix
        }

    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[progress * (this->tileLength)], this->tileLength);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyIn2(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1)))
        {
            if (progress == 0 || progress == 1)
            {
                if (this->lasttileLength >= (progress+1) * this->tileLength)
                {
                    DataCopy(xLocal, xGm[progress * (this->tileLength)],
                             (this->tileLength));
                }
                else
                {
                    DataCopy(xLocal, xGm[progress * (this->tileLength)],
                             (this->tileLength));
                    for (int32_t j = this->lasttileLength%this->tileLength; j < this->tileLength; j++)
                    {
                        xLocal.SetValue(j, this->spareNum);
                    }
                }
            }
            else
            {
                DataCopy(
                    xLocal,
                    xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                    (this->tileLength));
            }
        }
        else
        {
            DataCopy(xLocal, xGm[progress * (this->tileLength)],
                     (this->tileLength));
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        if constexpr (std::is_same_v<DTYPE_X, float>)
        {
            auto tmp3Local = Tmp3.Get<int32_t>();
            DTYPE_X min = -this->min;
            DTYPE_X width = this->bins / (this->max - this->min);
            DTYPE_X Tvalue;
            // Fix
            LocalTensor<DTYPE_X> tmp1Local = Tmp1.Get<DTYPE_X>();
            LocalTensor<DTYPE_X> zeroLocal = calcZero.Get<DTYPE_X>();
            LocalTensor<DTYPE_X> tmp4Local = Tmp4.Get<DTYPE_X>();
            LocalTensor<DTYPE_X> yLocal = calcTensor.Get<DTYPE_X>();
            auto tmp2Local = Tmp2.Get<uint8_t>();
            Muls(zeroLocal, zeroLocal, (DTYPE_X)0, this->tileLength);
            Adds(tmp4Local, zeroLocal, (DTYPE_X)this->max, this->tileLength);
            Compare(tmp2Local, xLocal, tmp4Local, CMPMODE::NE, this->tileLength);
            Select(zeroLocal, tmp2Local, zeroLocal, static_cast<DTYPE_X>(1),
                   SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
            ReduceSum(yLocal, zeroLocal, tmp1Local, this->tileLength);
            Tvalue = (float)yLocal.GetValue(0) + (float)yGm.GetValue(this->bins - 1);
            yGm.SetValue(this->bins - 1, (DTYPE_X)Tvalue);
            Adds(xLocal, xLocal, min, this->tileLength);
            Muls(xLocal, xLocal, width, this->tileLength);
            Cast(tmp3Local, xLocal, RoundMode::CAST_FLOOR, this->tileLength);
            Cast(xLocal, tmp3Local, RoundMode::CAST_NONE, this->tileLength);
            for (int32_t i = 0; i < this->bins; i++)
            {
                LocalTensor<DTYPE_X> yLocal = calcTensor.Get<DTYPE_X>();
                LocalTensor<DTYPE_X> tmp1Local = Tmp1.Get<DTYPE_X>();
                auto tmp2Local = Tmp2.Get<uint8_t>();
                LocalTensor<DTYPE_X> tmp4Local = Tmp4.Get<DTYPE_X>();
                LocalTensor<DTYPE_X> zeroLocal = calcZero.Get<DTYPE_X>();
                Muls(zeroLocal, zeroLocal, (DTYPE_X)0, this->tileLength); // 全0
                Adds(tmp4Local, zeroLocal, (DTYPE_X)i, this->tileLength);
                Compare(tmp2Local, xLocal, tmp4Local, CMPMODE::NE, this->tileLength);
                Select(zeroLocal, tmp2Local, zeroLocal, static_cast<DTYPE_X>(1),
                       SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
                ReduceSum(yLocal, zeroLocal, tmp1Local, this->tileLength);
                Tvalue = (float)yGm.GetValue(i) + (float)yLocal.GetValue(0);
                yGm.SetValue(i, (DTYPE_X)Tvalue);
            }
            inQueueX.FreeTensor(xLocal);
        }
        else
        {
            auto xLocal_f = calcXF.Get<float>();
            auto tmp3Local = Tmp3.Get<int32_t>();
            float min = -this->min;
            float width = this->bins / (this->max - this->min);
            float Tvalue;
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, this->tileLength);
            // Fix
            auto tmp1Local = Tmp1.Get<float>();
            auto zeroLocal = calcZero.Get<float>();
            auto tmp4Local = Tmp4.Get<float>();
            LocalTensor<float> yLocal = calcTensor.Get<float>();
            auto tmp2Local = Tmp2.Get<uint8_t>();
            Muls(zeroLocal, zeroLocal, (float)0, this->tileLength);
            Adds(tmp4Local, zeroLocal, (float)this->max, this->tileLength);
            Compare(tmp2Local, xLocal_f, tmp4Local, CMPMODE::NE, this->tileLength);
            Select(zeroLocal, tmp2Local, zeroLocal, static_cast<float>(1),
                   SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
            ReduceSum(yLocal, zeroLocal, tmp1Local, this->tileLength);
            Tvalue = (float)yLocal.GetValue(0) + (float)yGm.GetValue(this->bins - 1);
            yGm.SetValue(this->bins - 1, (DTYPE_X)Tvalue);
            Adds(xLocal_f, xLocal_f, min, this->tileLength);
            Muls(xLocal_f, xLocal_f, width, this->tileLength);
            Cast(tmp3Local, xLocal_f, RoundMode::CAST_FLOOR, this->tileLength);
            Cast(xLocal_f, tmp3Local, RoundMode::CAST_NONE, this->tileLength);
            for (int32_t i = 0; i < this->bins; i++)
            {
                LocalTensor<float> yLocal = calcTensor.Get<float>();
                auto tmp1Local = Tmp1.Get<float>();
                auto tmp2Local = Tmp2.Get<uint8_t>();
                auto tmp4Local = Tmp4.Get<float>();
                auto zeroLocal = calcZero.Get<float>();
                Muls(zeroLocal, zeroLocal, (float)0, this->tileLength); // 全0
                Adds(tmp4Local, zeroLocal, (float)i, this->tileLength);
                Compare(tmp2Local, xLocal_f, tmp4Local, CMPMODE::NE, this->tileLength);
                Select(zeroLocal, tmp2Local, zeroLocal, static_cast<float>(1),
                       SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
                ReduceSum(yLocal, zeroLocal, tmp1Local, this->tileLength);
                Tvalue = (float)yGm.GetValue(i) + (float)yLocal.GetValue(0);
                yGm.SetValue(i, (DTYPE_X)Tvalue);
            }
            inQueueX.FreeTensor(xLocal);
        }
    }

    __aicore__ inline void SetMM(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();

        if constexpr (std::is_same_v<DTYPE_X, float>)
        {
            LocalTensor<DTYPE_X> yLocal = calcTensor.Get<DTYPE_X>();
            LocalTensor<DTYPE_X> tmp1Local = Tmp1.Get<DTYPE_X>();
            ReduceMax(yLocal, xLocal, tmp1Local, this->tileLength);
            this->max = ((float)yLocal.GetValue(0) > this->max) ? yLocal.GetValue(0) : this->max;
            ReduceMin(yLocal, xLocal, tmp1Local, this->tileLength);
            this->min = ((float)yLocal.GetValue(0) < this->min) ? yLocal.GetValue(0) : this->min;
            inQueueX.FreeTensor(xLocal);
        }
        else
        {
            LocalTensor<float> yLocal = calcTensor.Get<float>();
            auto xLocal_f = calcXF.Get<float>();
            auto tmp1Local = Tmp1.Get<float>();
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, this->tileLength);
            ReduceMax(yLocal, xLocal_f, tmp1Local, this->tileLength);
            this->max = (yLocal.GetValue(0) > this->max) ? yLocal.GetValue(0) : this->max;
            ReduceMin(yLocal, xLocal_f, tmp1Local, this->tileLength);
            this->min = (yLocal.GetValue(0) < this->min) ? yLocal.GetValue(0) : this->min;
            inQueueX.FreeTensor(xLocal);
        }
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TBuf<QuePosition::VECCALC> Tmp1, Tmp2, Tmp3, Tmp4, calcZero, calcXF, calcTensor;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    int32_t blockLength;
    int32_t tileNum;
    int32_t tileLength;
    int32_t lasttileLength;
    int64_t bins;
    DTYPE_X spareNum;
    float min;
    float max;
    uint32_t usedataValue;
};

extern "C" __global__ __aicore__ void histogram(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelHistc op;

    op.Init(x, y, tiling_data.bins, tiling_data.min, tiling_data.max,
            tiling_data.blockLength, tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength);
    op.Process();
} 
